import random
import pandas as pd
import csv
import numpy as np
from utils.base_util import new_fold, get_list_from_dict
from utils.painter import draw_hist
from network.ResNet_model import resnet50
from network.Basic_model import basic_net
from network.metrics import max_distance
import torch
import utils.image_process as imp
from DRR.drr_generator import Projector
from configs.predict_conf import load_predict_config
from configs.config import global_var
from utils.label_transform import label2real
import matplotlib.pyplot as plt
from tqdm import tqdm


class Predict:
    def __init__(self, mode, num_classes, model_path, CT_path, saving_drr_path):
        self.mode = mode
        # self.model = resnet50(num_classes=num_classes)
        self.model = basic_net(out_channels=num_classes)
        self.model.load_state_dict(torch.load(model_path))
        # 导入用于DRR生成的对象
        self.data_generator = Projector(directory=CT_path)
        # 图片保存位置
        self.saving_drr_path = saving_drr_path

    def get_drr(self, pre_alpha, pre_beta, pre_theta, pre_tx, pre_ty, pre_tz):
        new_drr = imp.get_new_drr_numpy(pre_alpha, pre_beta, pre_theta,
                                        pre_tx, pre_ty, pre_tz, self.data_generator)
        return new_drr

    def predict(self, test_img, _tru_r, _tru_t, label):
        if self.mode == 'reg':
            self.model.eval()
            with torch.no_grad():
                test_img = test_img.reshape(1, 1, 128, 128) / 255
                img_tensor = torch.from_numpy(test_img).float()
                out, hidden_vector = self.model(img_tensor)
            pre_r = out[:, 0: 3].detach().cpu().numpy()
            pre_t = out[:, 3: 6].detach().cpu().numpy()
            # 获取预测的DRR
            pre_alpha, pre_beta, pre_theta, pre_tx, pre_ty, pre_tz = label2real(pre_r, pre_t, 1)
            # imp.get_DRR_image(projector=self.data_generator,
            #                   tru_img=test_img,
            #                   tru_r=_tru_r,
            #                   tru_t=_tru_t,
            #                   pre_r=[float(pre_alpha), float(pre_beta), float(pre_theta)],
            #                   pre_t=[float(pre_tx), float(pre_ty), float(pre_tz)],
            #                   i=label,
            #                   saving_path=self.saving_drr_path + "drr_vis_result")
            return np.array([float(pre_alpha), float(pre_beta), float(pre_theta)]), \
                np.array([float(pre_tx), float(pre_ty), float(pre_tz)]), hidden_vector.detach().cpu().numpy()


def dist_visualization(value_name, y1, y2, y3, idx, title, xy_label, path):
    _y1 = get_list_from_dict(y1, idx)
    _y2 = get_list_from_dict(y2, idx)
    _y3 = get_list_from_dict(y3, idx)
    draw_hist(value_name,
              [_y1, _y2, _y3],
              '平移{}:真实值-预测值1-预测值2'.format(title),
              xy_label,
              path)


if __name__ == "__main__":
    new_fold('pre_result', ['noise_result', 'statistical_result'])
    # 导入相应的模型
    pred_conf = load_predict_config(global_var.loading_address)
    predictor = Predict(**pred_conf)
    # 记录运行时的所有真实值，预测值
    tru = pd.DataFrame(columns=['tru_rx', 'tru_ry', 'tru_rz', 'tru_tx', 'tru_ty', 'tru_tz'])
    pre = pd.DataFrame(columns=['pre_rx', 'pre_ry', 'pre_rz', 'pre_tx', 'pre_ty', 'pre_tz'])
    hidden1 = pd.DataFrame(columns=[('h' + str(i)) for i in range(9248)])
    hidden2 = pd.DataFrame(columns=[('h' + str(i)) for i in range(9248)])
    VM = []
    # 每次只对一个维度加噪声
    noise = [0, 0, 5, 0, 0, 0]  # 前三维是角度对应的噪声，后三维是平移对应的噪声
    img_name = 'rz'     # 实际生成的对比DRR图片的标题
    for epoch in tqdm(range(1000)):
        # 随机生成一组位姿参数
        # 标准正位：（0±15， 270±5， 90±5）
        rx = np.random.uniform(-noise[0], noise[0])
        ry = np.random.uniform(270 - noise[1], 270 + noise[1])
        rz = np.random.uniform(90 - noise[2], 90 + noise[2])
        # 位移±50
        tx = np.random.uniform(-noise[3], noise[3])
        ty = np.random.uniform(-noise[4], noise[4])
        tz = np.random.uniform(-noise[5], noise[5])
        # 获取加噪后生成的图片
        test_drr = predictor.get_drr(rx, ry, rz, tx, ty, tz).reshape(128, 128, 1)
        # 通过网络获取预测的输出
        pre_r1, pre_t1, hidden_pre1 = predictor.predict(test_drr, [rx, ry, rz], [tx, ty, tz], epoch)
        # 根据网络的输出投影出对应的DRR
        pre_drr = predictor.get_drr(pre_r1[0], pre_r1[1], pre_r1[2],
                                    pre_t1[0], pre_t1[1], pre_t1[2]).reshape(128, 128, 1)
        # 将网络的预测值生成的drr再输入网络，获取其对应的潜空间的隐变量
        _, _, hidden_pre2 = predictor.predict(pre_drr, pre_r1, pre_t1, epoch)
        # 计算真实值和预测值之间的VM
        vm = max_distance(np.array([rx, ry, rz]), np.array([tx, ty, tz]),
                           pre_r1, pre_t1, 1, is_label_transform=False)
        MPD.append(vm)
        # 数据记录
        tru.loc[epoch] = [rx, ry, rz, tx, ty, tz]
        pre.loc[epoch] = [pre_r1[0], pre_r1[1], pre_r1[2], pre_t1[0], pre_t1[1], pre_t1[2]]
        hidden_pre1 = hidden_pre1.ravel().tolist()
        hidden_pre2 = hidden_pre2.ravel().tolist()
        hidden1.loc[epoch] = hidden_pre1
        hidden2.loc[epoch] = hidden_pre2
        # 可视化
        this_tru = np.array([rx, ry, rz, tx, ty, tz]).ravel()
        this_pre = np.array([pre_r1[0], pre_r1[1], pre_r1[2], pre_t1[0], pre_t1[1], pre_t1[2]]).ravel()
        plt.subplot(1, 2, 1)
        plt.title('实际值'
                  + "\n" +
                  '角度' + str(np.around(this_tru[0:3], 2))
                  + "\n" +
                  '平移' + str(np.around(this_tru[3:], 2)))
        plt.imshow(test_drr, cmap='gray')
        plt.subplot(1, 2, 2)
        plt.title('预测值'
                  + "\n" +
                  '角度' + str(np.around(this_pre[0:3], 2))
                  + "\n" +
                  '平移' + str(np.around(this_pre[3:], 2)))
        plt.imshow(pre_drr, cmap='gray')
        plt.suptitle("{0}:验证{1}加噪的精度".format(epoch, img_name))
        plt.savefig(predictor.saving_drr_path + "noise_result/第{0}组.png".format(epoch))
        # plt.show()
    # 写入csv文件
    tru.to_csv(predictor.saving_drr_path + "statistical_result/tru.csv")
    pre.to_csv(predictor.saving_drr_path + "statistical_result/pre.csv")
    hidden1.to_csv(predictor.saving_drr_path + "statistical_result/hidden1.csv")
    hidden2.to_csv(predictor.saving_drr_path + "statistical_result/hidden2.csv")
    with open(predictor.saving_drr_path + "statistical_result/MPD.csv", 'w') as mpd_file:
        writer = csv.writer(mpd_file)
        # 写入MPD
        writer.writerow(MPD)
    # # 将数据可视化
    # # 真实的值和预测值，以及扰动之后的预测值之间的对比
    # # tx
    # dist_visualization(y1=tru,
    #                    y2=pre1,
    #                    y3=pre2,
    #                    idx=0,
    #                    value_name=['tru_rx', 'pre1_rx', 'pre2_rx'],
    #                    title='rx',
    #                    xy_label=['角度', '°'],
    #                    path=predictor.saving_drr_path + 'statistical_result/' + '角度rx.html')
    # # ty
    # dist_visualization(y1=tru,
    #                    y2=pre1,
    #                    y3=pre2,
    #                    idx=1,
    #                    value_name=['tru_ry', 'pre1_ry', 'pre2_ry'],
    #                    title='ry',
    #                    xy_label=['角度', '°'],
    #                    path=predictor.saving_drr_path + 'statistical_result/' + '角度ry.html')
    # # tz
    # dist_visualization(y1=tru,
    #                    y2=pre1,
    #                    y3=pre2,
    #                    idx=2,
    #                    value_name=['tru_rz', 'pre1_rz', 'pre2_rz'],
    #                    title='rz',
    #                    xy_label=['角度', '°'],
    #                    path=predictor.saving_drr_path + 'statistical_result/' + '角度rz.html')
    # #
    # # tx-noise
    # dist_visualization(y1=tp_margin,
    #                    y2=pp_margin,
    #                    y3=noise,
    #                    idx=0,
    #                    value_name=['tp-x', 'pp-x', 'n-x'],
    #                    title='margin-rx',
    #                    xy_label=['角度', '°'],
    #                    path=predictor.saving_drr_path + 'statistical_result/' + '角度rx-margin.html')
    # # ty-noise
    # dist_visualization(y1=tp_margin,
    #                    y2=pp_margin,
    #                    y3=noise,
    #                    idx=1,
    #                    value_name=['tp-y', 'pp-y', 'n-y'],
    #                    title='margin-ry',
    #                    xy_label=['角度', '°'],
    #                    path=predictor.saving_drr_path + 'statistical_result/' + '角度ry-margin.html')
    # # ty-noise
    # dist_visualization(y1=tp_margin,
    #                    y2=pp_margin,
    #                    y3=noise,
    #                    idx=2,
    #                    value_name=['tp-z', 'pp-z', 'n-z'],
    #                    title='margin-rz',
    #                    xy_label=['角度', '°'],
    #                    path=predictor.saving_drr_path + 'statistical_result/' + '角度rz-margin.html')
    # # 验证noise与偏移量之间的关系
    # rx_margin = get_list_from_dict(pp_margin, 0)
    # ry_margin = get_list_from_dict(pp_margin, 1)
    # rz_margin = get_list_from_dict(pp_margin, 2)
    # noise = get_list_from_dict(noise, 6)
    # plt.close()
    # plt.figure(figsize=(12, 5))
    # # plt.gca().set_aspect(1)
    # plt.subplot(1, 3, 1)
    # plt.title('pp_rx-n')
    # plt.xlabel('noise')
    # plt.ylabel('rx_margin')
    # plt.scatter(noise, rx_margin)
    # plt.subplot(1, 3, 2)
    # plt.title('pp_ry-n')
    # plt.xlabel('noise')
    # plt.ylabel('ry_margin')
    # plt.scatter(noise, ry_margin)
    # plt.subplot(1, 3, 3)
    # plt.title('pp_rz-n')
    # plt.xlabel('noise')
    # plt.ylabel('rz_margin')
    # plt.scatter(noise, rz_margin)
    # plt.savefig(predictor.saving_drr_path + 'statistical_result/' + 'r_margin-noise.png')
    # plt.show()
